…-topk

  Replaces the CPU argmax path in _sample_token with a single greedy_sampling
  kernel that does RMSNorm + lm_head matmul + all_gather(full logits) + global
  topk entirely on device. Per-step DtoH is now (B,) uint32 = 4 bytes instead
  of (B, vocab_per_device) f32 = ~248 KB, and gloo all_gather + torch argmax
  on CPU are both eliminated.

  Kernel graph is gather-then-topk, not the usual topk-then-gather-then-index:
  neuronx-cc 2.23 miscompiles topk when all_gather is downstream in the same
  kernel for certain vocab_per_device sizes (including 62080 used at TP=4),
  producing wrong token IDs. Keeping topk strictly downstream of all_gather
  sidesteps this, and taking topk over the gathered full-vocab logits directly
  returns the global winner ID -- no rank-offset arithmetic, no dynamic index.

  - kernels/sampling.py: replace compute_logits with greedy_sampling
  - qwen3_5.py: compile greedy_sampling, swap (B, vpd) f32 logits buffers
    for (B,) uint32 next_id buffers, simplify _sample_token to a single
    DtoH + reshape

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>